import torch
import torch.nn as nn


class SingleHead(nn.Module):
    def __init__(self, in_channel, out_channel, bias_fill=False, bias_value=0):
        super(SingleHead, self).__init__()
        self.feat_conv = nn.Conv2d(in_channel, in_channel, kernel_size=3, padding=1)
        self.relu = nn.ReLU()
        self.out_conv = nn.Conv2d(in_channel, out_channel, kernel_size=1)
        if bias_fill:
            self.out_conv.bias.data.fill_(bias_value)

    def forward(self, x):
        x = self.feat_conv(x)
        x = self.relu(x)
        x = self.out_conv(x)
        return x


class CenternetHead(nn.Module):
    """
    The head used in CenterNet for object classification and box regression.
    It has three subnet, with a common structure but separate parameters.
    """
    def __init__(self, head_conv=64, num_classes=1, landmark=None):
        super(CenternetHead, self).__init__()
        self.cls_head = SingleHead(
            head_conv,
            num_classes,
            bias_fill=True,
            bias_value=-2.19,
        )
        self.ltrb_head = SingleHead(head_conv, 4)
        self.landmark = landmark
        if landmark is not None:
            self.landmark_head = SingleHead(head_conv, landmark)



    def forward(self, x):
        cls = self.cls_head(x).sigmoid()
        ltrb = self.ltrb_head(x).exp()
        
        pred = {
            'cls': cls,
            'ltrb': ltrb,
        }
        if self.landmark is not None:
            pred['landm'] = self.landmark_head(x)
        return pred